import os
import cv2
import numpy as np
import torch

from lib.models.bisenetv2 import BiSeNetV2 as _BiSeNetV2
from utils.data_utils import preprocess_image, postprocess_output


class RoadSegmentor:
    def __init__(self, weights_path=None, device=None, num_classes=19):
        here = os.path.dirname(os.path.abspath(__file__))
        self.weights_path = weights_path or os.path.join(here, 'models', 'weights', 'model_final.pth')
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        self.num_classes = num_classes

        # 跳过模型内部的预训练下载行为
        _BiSeNetV2.load_pretrain = lambda self: None
        self.model = _BiSeNetV2(n_classes=self.num_classes, aux_mode='eval')
        self.model.to(self.device)
        self.model.eval()

        self._load_weights()

        # 路面类别 id（Cityscapes 常见为 0）
        self.road_class_id = 0

    def _load_weights(self):
        if not os.path.isfile(self.weights_path):
            print(f"[warn] 未找到权重: {self.weights_path}")
            print("       请将 model_final.pth 放入 deploy/models/weights/ 或运行 get_weights.py")
            return
        try:
            state = torch.load(self.weights_path, map_location=self.device)
            if 'state_dict' in state:
                state = state['state_dict']
            # 兼容不同保存格式
            self.model.load_state_dict(state, strict=False)
            print(f"[ok] 已加载权重: {self.weights_path}")
        except Exception as e:
            print(f"[error] 加载权重失败: {e}")

    def infer(self, bgr_frame: np.ndarray, return_overlay: bool = True):
        h, w = bgr_frame.shape[:2]
        inp = preprocess_image(bgr_frame)
        inp = inp.to(self.device)
        with torch.no_grad():
            out_tuple = self.model(inp)
            logits = out_tuple[0] if isinstance(out_tuple, (list, tuple)) else out_tuple
        pred = postprocess_output(logits, (h, w))

        road_mask = (pred == self.road_class_id).astype(np.uint8)

        overlay = None
        if return_overlay:
            overlay = self.colorize_mask(road_mask, bgr_frame)
        return road_mask, overlay

    def colorize_mask(self, road_mask: np.ndarray, bgr_frame: np.ndarray):
        color = np.array([0, 200, 0], dtype=np.uint8)  # BGR 绿色
        mask_color = np.zeros_like(bgr_frame)
        mask_color[road_mask > 0] = color
        overlay = cv2.addWeighted(bgr_frame, 1.0, mask_color, 0.4, 0)
        return overlay